import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import math
import copy
from torch.utils.data import Dataset, DataLoader

# -------------------------- 1. 基础配置 --------------------------
seed = 42
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

Device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using Device: {Device}")


# -------------------------- 2. 数据集（不变） --------------------------
class SimpleDataset(Dataset):
    def __init__(self, num):
        prefix = torch.tensor([0])  # 起始符
        suffix = torch.tensor([9])  # 结束符
        self.data = [torch.randint(1, 9, size=(5,), dtype=torch.long) for _ in range(num)]
        self.labels = [torch.cat([prefix, _, suffix], dim=0) for _ in copy.deepcopy(self.data)]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]  # data: (5,), label: (7,)


def create_data(batch_size=32, num=1000):
    dataset = SimpleDataset(num)
    return DataLoader(dataset, batch_size=batch_size, shuffle=True)


# -------------------------- 3. 修复的 MultiHeadAttention（核心修改） --------------------------
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads  # 必须能整除（如512//8=64）

        # Q/K/V投影层 + 输出投影层
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)
        self.scale = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(Device)

    def forward(self, q, k, v, attn_mask=None, key_padding_mask=None):
        """
        修复维度转换逻辑：删除多余的 transpose(0,1)，确保批量维度一致
        输入维度：q/k/v → (seq_len, batch_size, d_model)
        输出维度：attn_output → (seq_len_q, batch_size, d_model)
        """
        seq_len_q, batch_size, _ = q.shape  # 如：(6, 32, 512)（解码器输入序列长度6，批次32）
        seq_len_k, _, _ = k.shape           # 如：(5, 32, 512)（编码器输出序列长度5）

        # -------------------------- 步骤1：Q/K/V 投影 + 拆分多头 --------------------------
        # 投影：(seq_len, batch_size, d_model) → 保持维度不变
        q_proj = self.w_q(q)  # (seq_len_q, batch_size, d_model)
        k_proj = self.w_k(k)  # (seq_len_k, batch_size, d_model)
        v_proj = self.w_v(v)  # (seq_len_k, batch_size, d_model)

        # 拆分多头：将 d_model 拆分为 n_heads × head_dim → (seq_len, batch_size, n_heads, head_dim)
        q_split = q_proj.view(seq_len_q, batch_size, self.n_heads, self.head_dim)  # 如：(6,32,8,64)
        k_split = k_proj.view(seq_len_k, batch_size, self.n_heads, self.head_dim)  # 如：(5,32,8,64)
        v_split = v_proj.view(seq_len_k, batch_size, self.n_heads, self.head_dim)  # 如：(5,32,8,64)

        # 转置 + 合并批量与多头维度：(seq_len, batch_size×n_heads, head_dim)
        # 目的：用 batch 维度承载多头，实现并行计算（无需循环每个头）
        q_trans = q_split.transpose(1, 2).contiguous()  # (seq_len_q, n_heads, batch_size, head_dim)
        q_trans = q_trans.view(seq_len_q, batch_size * self.n_heads, self.head_dim)  # 如：(6, 32×8=256, 64)
        k_trans = k_split.transpose(1, 2).contiguous().view(seq_len_k, batch_size * self.n_heads, self.head_dim)  # (5,256,64)
        v_trans = v_split.transpose(1, 2).contiguous().view(seq_len_k, batch_size * self.n_heads, self.head_dim)  # (5,256,64)

        # -------------------------- 步骤2：计算注意力分数（核心修复） --------------------------
        # 转置 q_trans 为 (batch_size×n_heads, seq_len_q, head_dim) → 适配 bmm 输入要求
        q_trans_for_bmm = q_trans.transpose(0, 1)  # 如：(256, 6, 64)
        # 转置 k_trans 为 (batch_size×n_heads, head_dim, seq_len_k) → 适配 bmm 输入要求
        k_trans_for_bmm = k_trans.transpose(0, 1).transpose(1, 2)  # 如：(256, 64, 5)

        # 批量矩阵乘法：(B, Lq, Hd) × (B, Hd, Lk) → (B, Lq, Lk)（B=batch×heads，Lq=seq_len_q，Lk=seq_len_k）
        attn_scores = torch.bmm(q_trans_for_bmm, k_trans_for_bmm)  # 如：(256, 6, 5)
        # 修复：删除多余的 transpose(0,1)！保持批量维度在前，后续计算一致

        # 缩放：除以 sqrt(head_dim) → 防止分数过大导致 softmax 梯度消失
        attn_scores = attn_scores / self.scale

        # -------------------------- 步骤3：应用掩码（维度匹配） --------------------------
        # 3.1 应用 attn_mask（如解码器下三角掩码，维度：(seq_len_q, seq_len_k)）
        if attn_mask is not None:
            # 扩展掩码维度：(Lq, Lk) → (1, Lq, Lk) → 广播到 (B, Lq, Lk)
            attn_mask = attn_mask.unsqueeze(0)  # 如：(1,6,5)
            attn_scores = attn_scores.masked_fill(attn_mask == 0, float('-1e9'))  # 掩码位置设为极小值

        # 3.2 应用 key_padding_mask（填充掩码，维度：(batch_size, seq_len_k)）
        if key_padding_mask is not None:
            # 扩展掩码维度：(B0, Lk) → (B0, 1, Lk) → 重复 n_heads 次 → (B0×n_heads, 1, Lk)
            key_padding_mask = key_padding_mask.unsqueeze(1).repeat(1, self.n_heads, 1)  # 如：(32,8,5)
            key_padding_mask = key_padding_mask.view(batch_size * self.n_heads, 1, seq_len_k)  # 如：(256,1,5)
            attn_scores = attn_scores.masked_fill(key_padding_mask == 1, float('-1e9'))

        # -------------------------- 步骤4：Softmax + Dropout + 乘 V --------------------------
        # Softmax：按最后一维（seq_len_k）计算 → 注意力权重（维度：(B, Lq, Lk)）
        attn_weights = F.softmax(attn_scores, dim=-1)  # 如：(256,6,5)
        # Dropout：防止过拟合
        attn_weights = self.dropout(attn_weights)

        # 转置 V 为 (B, Lk, Hd) → 适配 bmm 输入要求
        v_trans_for_bmm = v_trans.transpose(0, 1)  # 如：(256,5,64)
        # 批量矩阵乘法：(B, Lq, Lk) × (B, Lk, Hd) → (B, Lq, Hd)
        attn_output = torch.bmm(attn_weights, v_trans_for_bmm)  # 如：(256,6,64)

        # -------------------------- 步骤5：合并多头 + 输出投影 --------------------------
        # 转置为 (Lq, B, Hd) → 拆分批量与多头维度
        attn_output = attn_output.transpose(0, 1)  # 如：(6,256,64)
        # 合并多头：(Lq, B0×n_heads, Hd) → (Lq, B0, n_heads×Hd) = (Lq, B0, d_model)
        attn_output = attn_output.contiguous().view(seq_len_q, batch_size, self.d_model)  # 如：(6,32,512)

        # 输出投影：映射回 d_model 维度（可选，确保输出维度与输入一致）
        attn_output = self.w_o(attn_output)

        return attn_output, attn_weights


# -------------------------- 4. DecoderLayer（不变） --------------------------
class DecoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, dim_feedforward=2048, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.multihead_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = F.relu

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        tgt2, _ = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2, _ = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt


# -------------------------- 5. CustomTransformerDecoder（不变） --------------------------
class CustomTransformerDecoder(nn.Module):
    def __init__(self, decoder_layer, num_layers, norm=None):
        super(CustomTransformerDecoder, self).__init__()
        self.layers = nn.ModuleList([copy.deepcopy(decoder_layer) for _ in range(num_layers)])
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        output = tgt
        for layer in self.layers:
            output = layer(output, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
                           tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
        if self.norm is not None:
            output = self.norm(output)
        return output


# -------------------------- 6. 位置编码（不变） --------------------------
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=10, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :, :]
        return self.dropout(x)


# -------------------------- 7. EncoderLayer（不变） --------------------------
class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, dim_feedforward=2048, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.activation = F.relu

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        src2, _ = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src


# -------------------------- 8. CustomTransformerEncoder（不变） --------------------------
class CustomTransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers, norm=None):
        super(CustomTransformerEncoder, self).__init__()
        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)])
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src, mask=None, src_key_padding_mask=None):
        output = src
        for layer in self.layers:
            output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
        if self.norm is not None:
            output = self.norm(output)
        return output


# -------------------------- 9. Transformer 模型（不变） --------------------------
class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, seq_len, d_model=512,
                 dim_feedforward=2048, num_layers=6, n_heads=8, dropout=0.1):
        super(Transformer, self).__init__()
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.d_model = d_model
        self.pos_encoder = PositionalEncoding(d_model, max_len=seq_len, dropout=dropout)
        encoder_layer = EncoderLayer(d_model, n_heads, dim_feedforward, dropout)
        self.encoder = CustomTransformerEncoder(encoder_layer, num_layers)
        decoder_layer = DecoderLayer(d_model, n_heads, dim_feedforward, dropout)
        self.decoder = CustomTransformerDecoder(decoder_layer, num_layers)
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
        self.init_weights()

    def init_weights(self):
        initrange = 0.1
        self.src_embedding.weight.data.uniform_(-initrange, initrange)
        self.tgt_embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_()
        self.fc_out.weight.data.uniform_(-initrange, initrange)
        for m in self.modules():
            if isinstance(m, MultiHeadAttention):
                nn.init.xavier_uniform_(m.w_q.weight)
                nn.init.xavier_uniform_(m.w_k.weight)
                nn.init.xavier_uniform_(m.w_v.weight)
                nn.init.xavier_uniform_(m.w_o.weight)

    def generate_tgt_mask(self, tgt_seq_len):
        mask = torch.tril(torch.ones(tgt_seq_len, tgt_seq_len, dtype=torch.bool)).to(Device)
        return mask

    def forward(self, src, tgt):
        src_seq_len = src.size(1)
        tgt_seq_len = tgt.size(1)
        src_emb = self.src_embedding(src).transpose(0, 1) * math.sqrt(self.d_model)
        src_emb = self.pos_encoder(src_emb)
        tgt_emb = self.tgt_embedding(tgt).transpose(0, 1) * math.sqrt(self.d_model)
        tgt_emb = self.pos_encoder(tgt_emb)
        tgt_mask = self.generate_tgt_mask(tgt_seq_len)
        memory = self.encoder(src_emb, src_key_padding_mask=None)
        tgt_out = self.decoder(tgt_emb, memory, tgt_mask=tgt_mask,
                               tgt_key_padding_mask=None, memory_key_padding_mask=None)
        output = self.fc_out(tgt_out).transpose(0, 1)
        return output


# -------------------------- 10. 训练函数（不变） --------------------------
def train():
    datas = create_data(batch_size=32, num=1000)
    model = Transformer(
        src_vocab_size=11,
        tgt_vocab_size=11,
        seq_len=10,
        d_model=512,
        dim_feedforward=4 * 512,
        num_layers=6,
        n_heads=8
    ).to(Device)

    intervals = 100
    criterion = nn.CrossEntropyLoss()
    opt = optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-9)
    # opt = optim.SGD(model.parameters(), lr=0.001)

    for interval in range(intervals):
        model.train()
        loss_all = 0.0
        for data, label in datas:
            data = data.to(Device)
            label = label.to(Device)
            input_label = label[:, :-1]
            output_label = label[:, 1:]

            opt.zero_grad()
            output = model(data, input_label)
            loss = criterion(output.reshape(-1, output.size(-1)), output_label.reshape(-1))
            loss_all += loss.item()

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            opt.step()

        avg_loss = loss_all / len(datas)
        print(f'第 {interval + 1} 轮, 平均损失: {avg_loss:.6f}, 学习率: {opt.param_groups[0]["lr"]:.8f}')

    # 测试
    with torch.no_grad():
        test_in = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.long).to(Device)
        tgt_init = torch.tensor([[0]], dtype=torch.long).to(Device)
        print("\n测试序列生成：")
        while True:
            output = model(test_in, tgt_init)
            next_token = output.argmax(dim=-1)[:, -1].unsqueeze(1)
            tgt_init = torch.cat([tgt_init, next_token], dim=1)
            print(f"当前预测: {tgt_init.squeeze().tolist()}")
            if next_token.item() == 9 or tgt_init.size(1) > 10:
                break


# -------------------------- 11. 启动训练 --------------------------
if __name__ == "__main__":
    train()